#leetcode题目215：数组中的第K个最大元素
#难度：中等
#时间复杂度：O(nlogn)
#空间复杂度：O(1)
#方法：堆

# 注意：要求使用的是o（n）的时间复杂度，所以不能使用堆排序，而是快排？你必须设计并实现时间复杂度为 O(n) 的算法解决此问题。

from typing import List
import random

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        """
        数组中的第K个最大元素
        给定整数数组 nums 和整数 k，请返回数组中第 k 个最大的元素。
        请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。
        你必须设计并实现时间复杂度为 O(n) 的算法解决此问题。
        """
        n = len(nums)
        l = 0
        r = n - 1
        k = n - k  # 转换为第k个最小元素
        
        while l <= r:
            # 分区操作
            range_result = self.partition(nums, l, r)
            
            if range_result[0] <= k <= range_result[1]:
                # 如果k在等于pivot的范围内，找到了答案
                return nums[range_result[0]]
            elif range_result[0] > k:
                # 如果k在左半部分
                r = range_result[0] - 1
            else:  # range_result[1] < k
                # 如果k在右半部分
                l = range_result[1] + 1
        
        return 0  # 理论上不会到达这里

    def partition(self, nums: List[int], l: int, r: int) -> List[int]:
        """
        三路分区函数
        将数组分为三部分：小于pivot、等于pivot、大于pivot
        返回等于pivot的元素的起始和结束索引
        """
        # 随机选择pivot
        choose_index = l + random.randint(0, r - l)
        target = nums[choose_index]
        
        index = l
        start = l - 1  # 小于target的右边界
        end = r        # 大于target的左边界
        
        while index <= end:
            if nums[index] > target:
                # 大于target，交换到右边
                self.swap(nums, end, index)
                end -= 1
            elif nums[index] < target:
                # 小于target，交换到左边
                start += 1
                self.swap(nums, start, index)
                index += 1
            else:
                # 等于target，保持不变
                index += 1
        
        return [start + 1, end]
    
    def swap(self, nums: List[int], a: int, b: int) -> None:
        """
        交换数组中两个位置的元素
        """
        nums[a], nums[b] = nums[b], nums[a]


#测试数据
nums = [3,2,1,5,6,4]
k = 2
#预期输出：5
solution = Solution()
print(solution.findKthLargest(nums, k))

nums = [3,2,3,1,2,4,5,5,6]
k = 4
#预期输出：4
solution = Solution()
print(solution.findKthLargest(nums, k))